# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: yuxuan

import datetime
import functools
import logging
import signal

from absl import (flags, app)
from tornado.ioloop import IOLoop
import numpy as np

import coin.flow.subscriber as flow_subscriber
from coin.exchange.kr_rest.product_holders import get_holder_from_product
from coin.exchange.base.order_gateway import OrderSide, Order
from coin.strategy.hk_hard_arb.products import NativeProducts
from coin.exchange.base.kr_rest.product import NativeProduct
from coin.base.book.book_builder_base import BookBuilderBase
from coin.base.param_util import to_list
from coin.exchange.shared.kr_rest.product import SharedProduct
from coin.exchange.util.feed_checker import get_default_feed_checker_config
from coin.flow.simple_queue import SimpleQueue
from coin.proto.coin_order_checker_pb2 import AccountCheckerConfig
from coin.strategy.executor.util.cancel_all_executor import (SimpleCancelAllExecutor)
from coin.strategy.basis_strat.balance_printer import KimchiEstimatedBalancePrinter
from coin.strategy.hk_hard_arb.util import load_clean_json
from coin.strategy.mm.base.multi_feed_subsystem import MultiFeedSubsystem
from coin.strategy.mm.feed import SubscriptionRequest, get_flow_book_builder_cls
from coin.strategy.mm.live_envs import (create_live_env_instance, get_live_env_cls)
from coin.strategy.executor.aggressive.factory import construct_agg_executor
from coin.strategy.executor.passive.factory import construct_mm_executor
from coin.strategy.executor.passive.util import \
    convert_executor_config_from_dict as convert_pass_config
from coin.strategy.executor.aggressive.util import \
    convert_executor_config_from_dict as convert_agg_config
from coin.support.accounting.checker.account_checker import AccountChecker
from coin.strategy.accounting.pnl_balance_printer.util import reserve_as_proto
from coin.support.accounting.checker.config_util import convert_checker_config_from_dict
from coin.exchange.base.strategy_logger import StrategyLogProvider
from coin.exchange.base.order_gateway_config import convert_og_config_from_dict
from coin.strategy.util.request_util import convert_strategy_request_from_dict
from coin.strategy.basis_strat.fill_info_writer import FillInfo, FillInfoWriter
from coin.strategy.accounting.fee import get_fee_rate

flags.DEFINE_string('og_configs', '', '')

flags.DEFINE_string('config_filename', '', '')

flags.DEFINE_integer('exit_after_min', 30, '')

flags.DEFINE_boolean('dry_run', False, '')

flags.DEFINE_string('trading_tickers', '', '')

flags.DEFINE_boolean('guess_fill_from_balance', True, '')

flags.DEFINE_boolean('guess_fill_from_balance_only', False, '')

flags.DEFINE_boolean('guess_fill_from_book', False, 'Not used for now.')

flags.DEFINE_float('aggressive_bp', 3, '')

flags.DEFINE_boolean('use_stop_qty_pricer', False, 'By default, use BestPricePricer.')

flags.DEFINE_float('book_invalid_sec', 5, '')

flags.DEFINE_boolean('update_pass_pos_with_balance', False, '')

flags.DEFINE_boolean('update_agg_pos_with_balance', False, '')

flags.DEFINE_float('update_pos_with_balance_sec', 4, '')

flags.DEFINE_integer('post_pass_order_within_levels', 0, '0 turns off this feature.')

flags.DEFINE_float('post_pass_order_within_lotsize_multiplier', 0.0, '0 to turn off this feature.')

flags.DEFINE_string('strategy_name', '', '')

flags.DEFINE_boolean('dump_pnl_balance', True, '')

flags.DEFINE_integer('book_levels', 5, '')

flags.DEFINE_boolean('subscribe_trade', False, '')

flags.DEFINE_float('max_pos_lean_bp', 3, '')

FLAGS = flags.FLAGS


def _get_ticker(obj):
  if isinstance(obj, NativeProduct):
    ticker = obj.base.currency
  elif isinstance(obj, BookBuilderBase):
    ticker = obj.norm_symbol.split('-')[0]
  elif isinstance(obj, Order):
    ticker = _get_ticker(obj.product)
  else:
    assert False, type(obj)
  return ticker


def _gen_arb_products(common, products0, products1, ProductType0, ProductType1):
  if 'exch0_quote_currency' in common and 'exch1_quote_currency' in common:
    exch0_quote_currency = common['exch0_quote_currency']
    exch1_quote_currency = common['exch1_quote_currency']
  else:
    raise ValueError(common)

  products0 = NativeProducts.FromTypedStrProduct(exch0_quote_currency,
                                                 sorted(products0.keys()),
                                                 ProductType0)
  products1 = NativeProducts.FromTypedStrProduct(exch1_quote_currency,
                                                 sorted(products1.keys()),
                                                 ProductType1)
  assert len(products0.products) == len(products1.products)
  return products0, products1


class StopQtyPricer(object):
  @staticmethod
  def _find_price_with_qty(bids, asks, qty):
    buy_price = asks[-1][0]
    acc_qty = 0
    for ask in asks:
      acc_qty += ask[1]
      if acc_qty >= qty:
        buy_price = ask[0]
        break
    sell_price = bids[-1][0]
    acc_qty = 0
    for bid in bids:
      acc_qty += bid[1]
      if acc_qty >= qty:
        sell_price = bid[0]
        break
    return buy_price, sell_price

  @staticmethod
  def calculate_price(bids0, asks0, bids1, asks1, mean_scale, spread_std, threshold, lot_size):
    mid0 = SpreadBasisStrategy.get_mid(bids0, asks0)
    mid1 = SpreadBasisStrategy.get_mid(bids1, asks1) * mean_scale
    mid = (mid0 + mid1) / 2.0
    # Aggression at product1
    buy_stop_price, sell_stop_price = StopQtyPricer._find_price_with_qty(bids1, asks1, lot_size)
    assert buy_stop_price > sell_stop_price
    buy_price_base = buy_stop_price * mean_scale
    pass_sell_price0 = (buy_price_base + mid * threshold * spread_std)
    sell_price_base = sell_stop_price * mean_scale
    pass_buy_price0 = (sell_price_base - mid * threshold * spread_std)
    return float(pass_buy_price0), float(pass_sell_price0)

  @staticmethod
  def calculate_qty(bids1, asks1, pass_bal, max_pos, min_pos, min_pass_order_qty, lot_size):
    buy_qty = min(lot_size, max((max_pos - pass_bal) * 0.95, 0))
    sell_qty = min(lot_size, (pass_bal - min_pos) * 0.95)
    if buy_qty <= min_pass_order_qty:
      buy_qty = None
    if sell_qty <= min_pass_order_qty:
      sell_qty = None
    if buy_qty and sell_qty:
      if sell_qty > 2.99 * buy_qty:
        buy_qty = None
      elif buy_qty > 2.99 * sell_qty:
        sell_qty = None
    return buy_qty, sell_qty


class BestPricePricer(object):
  @staticmethod
  def calculate_price(bids0, asks0, bids1, asks1, mean_scale, spread_std, threshold, lot_size):
    mid0 = SpreadBasisStrategy.get_mid(bids0, asks0)
    mid1 = SpreadBasisStrategy.get_mid(bids1, asks1) * mean_scale
    mid = (mid0 + mid1) / 2.0
    # Aggression at product1
    buy_price_base = asks1[0][0] * mean_scale
    pass_sell_price0 = (buy_price_base + mid * threshold * spread_std)
    sell_price_base = bids1[0][0] * mean_scale
    pass_buy_price0 = (sell_price_base - mid * threshold * spread_std)
    return float(pass_buy_price0), float(pass_sell_price0)

  @staticmethod
  def calculate_qty(bids1, asks1, pass_bal, max_pos, min_pos, min_pass_order_qty, lot_size):
    buy_qty = min(lot_size, max((max_pos - pass_bal) * 0.95, 0))
    sell_qty = min(lot_size, (pass_bal - min_pos) * 0.95)
    if buy_qty <= min_pass_order_qty:
      buy_qty = None
    if sell_qty <= min_pass_order_qty:
      sell_qty = None
    if buy_qty is not None:
      buy_qty = min(buy_qty, asks1[0][1] * 0.6)
    if sell_qty is not None:
      sell_qty = min(sell_qty, bids1[0][1] * 0.6)
    if buy_qty and sell_qty:
      if sell_qty > 2.99 * buy_qty:
        buy_qty = None
      elif buy_qty > 2.99 * sell_qty:
        sell_qty = None
    return buy_qty, sell_qty


class SingleSpreadBasisSubStrategy(object):
  def __init__(self,
               parent,
               product0,
               product1,
               pull_edge,
               push_edge,
               trading_threshold=2,
               fillinfo_writer=None):
    self._product0 = product0
    self._product1 = product1
    self._ticker0 = _get_ticker(self._product0)
    self._ticker1 = _get_ticker(self._product1)
    self._tick_info0 = (get_holder_from_product(product0).product_info.gen_tick(),
                        get_holder_from_product(product0).product_info.gen_tick_qty())
    self._tick_info1 = (get_holder_from_product(product1).product_info.gen_tick(),
                        get_holder_from_product(product1).product_info.gen_tick_qty())
    self._config0 = None
    self._config1 = None
    self.og0 = None
    self.og1 = None
    self._parent = parent
    self.logger = parent.logger
    self.feed = parent._last_feed
    self._pull_edge = pull_edge * 1e-4
    self._push_edge = push_edge * 1e-4
    self._threshold = trading_threshold
    self._pass_executor0 = None
    self._agg_executor1 = None
    self._fillinfo_writer = fillinfo_writer
    self._pos0 = 0
    self._pos1 = 0
    self._pos0_ts = 0
    self._pos1_ts = 0
    self._fill0_ts = 0
    self._og_pos0 = 0
    self._og_pos1 = 0
    self._og_pos0_ts = 0
    self._og_pos1_ts = 0
    self._update_pass_pos_with_balance = FLAGS.update_pass_pos_with_balance
    self._update_agg_pos_with_balance = FLAGS.update_agg_pos_with_balance
    self._update_position_period = FLAGS.update_pos_with_balance_sec * 10**9
    self._spread_std = None
    self._aggressive_amt = FLAGS.aggressive_bp * 1e-4
    self._book_invalid_period = FLAGS.book_invalid_sec * 10**9
    self._book_ts = [0.0, 0.0]
    if FLAGS.use_stop_qty_pricer:
      self._pricer = StopQtyPricer
    else:
      self._pricer = BestPricePricer
    # Configs
    self._max_pos = None
    self._min_pos = None
    self._lot_size = None
    self._min_lot_size_ratio = None
    self._min_pass_order_qty = None
    self._pass_maker_fee = None
    self._agg_taker_fee = None
    # Handle order fill limit.
    self._post_pass_within_levels = FLAGS.post_pass_order_within_levels
    self._post_pass_within_lotsize = FLAGS.post_pass_order_within_lotsize_multiplier
    self._check_post_within = self._post_pass_within_levels > 0 \
        or self._post_pass_within_lotsize > 0.01
    self._last_pass_buy_price = 0
    self._last_pass_sell_price = 1e10
    # Handle fill message or balance update delay.
    self._guess_fill_from_balance = FLAGS.guess_fill_from_balance
    self._guess_fill_from_balance_only = FLAGS.guess_fill_from_balance_only
    if self._guess_fill_from_balance_only:
      assert self._guess_fill_from_balance
    self._guess_fill_from_book = FLAGS.guess_fill_from_book
    self._expected_agg_pos = None
    self._latest_pass_pos = None
    self._latest_pass_pos_from_fill = None
    self._equal_theshold = None
    self._expected_working_qty = 0
    self._last_agg_submit_fill_ts = 0
    # Handle agg working order update delay.
    self._agg_working_order_fills = dict()
    # Prevent feed invalid log explosion.
    self._last_feed_invalid_log_ts = 0

    self._config0_proto = None
    self._config1_proto = None

  def _is_book_valid(self):
    return self._get_ts() - min(self._book_ts) < self._book_invalid_period

  def _get_ts(self):
    return self._parent._ts

  def _is_pass_pos_valid(self):
    ts = self._get_ts()
    return ts - max(self._pos0_ts, self._og_pos0_ts, self._fill0_ts) \
        > self._update_position_period

  def _is_agg_pos_valid(self):
    ts = self._get_ts()
    return ts - max(self._pos1_ts, self._og_pos1_ts) > self._update_position_period

  def _get_pass_balance(self):
    og_pos = self.og0.get_balance(self._ticker0)
    if self._update_pass_pos_with_balance:
      if not self._position_equal(og_pos, self._og_pos0):
        self._og_pos0 = og_pos
        self._og_pos0_ts = self._get_ts()
      elif not self._position_equal(self._pos0, self._og_pos0) \
              and self._is_pass_pos_valid():
        self.logger.info('Update pos0, latest_pass_pos_from_fill with og_pos0, %f %f %f' %
                         (self._pos0, self._latest_pass_pos_from_fill, self._og_pos0))
        self._update_pos0(self._og_pos0)
        self._latest_pass_pos = self._latest_pass_pos_from_fill = self._og_pos0
        self._update_expected_agg_pos()
    else:
      if self._position_equal(self._pos0, og_pos):
        self._pos0 = og_pos
    return self._pos0

  def _get_agg_balance(self):
    if self._update_agg_pos_with_balance:
      self._og_pos1 = self.og1.get_balance(self._ticker1)
      if self._get_ts() - self._last_agg_submit_fill_ts > \
              self._update_position_period and \
              not self._position_equal(self._og_pos1, self._pos1):
        self.logger.info('Update pos1 with og_pos1, %f %f' % (self._pos1, self._og_pos1))
        self._pos1 = self._og_pos1
    return self._pos1

  def _update_pos0(self, pos):
    self._pos0 = pos
    self._pos0_ts = self._get_ts()

  def _update_pos1(self, pos):
    self._pos1 = pos
    self._pos1_ts = self._get_ts()

  def write_executor_config(self):
    if self._config0_proto is not None:
      self._strategy_logger.write_executor_config(self._config0_proto)
    if self._config1_proto is not None:
      self._strategy_logger.write_executor_config(self._config1_proto)

  def init_executors(self, req0, req1, config0, config1, og0, og1):
    self.og0 = og0
    self.og1 = og1
    self._config0 = config0
    self._config0['force_fire'] = True
    self._config1 = config1.copy()
    self._config1['force_fire'] = True
    self._config0_proto = convert_pass_config(self._product0.symbol, config0)
    self._config1_proto = convert_agg_config(self._product1.symbol, self._config1)
    self._pass_executor0 = construct_mm_executor(req0, self._config0_proto, self.og0, self.logger)
    self._agg_executor1 = construct_agg_executor(req1, self._config1_proto, self.og1, self.logger)
    # Fill smaller than this value will cause serious problem.
    self._equal_theshold = max(self._config0['lot_size'] * 1e-4, 1e-5)
    self._update_pos0(self.og0.get_balance(self._ticker0))
    self._update_pos1(self.og1.get_balance(self._ticker1))
    self._og_pos0 = self._pos0
    self._og_pos1 = self._pos1
    self._og_pos0_ts = self._og_pos1_ts = self._get_ts()
    self._expected_agg_pos = self._config0['max_pos'] - self._get_pass_balance()
    self._latest_pass_pos_from_fill = self._latest_pass_pos = self._pos0
    # Init configs
    self._max_pos = self._config0['max_pos']
    self._min_pos = self._config0['min_pos']
    self._lot_size = self._config0['lot_size']
    self._min_lot_size_ratio = self._config0['min_lot_size_ratio']
    self._min_pass_order_qty = self._lot_size * self._min_lot_size_ratio
    self._pass_maker_fee, _ = get_fee_rate(self.og0.market_type, self.og0.exchange)
    _, self._agg_taker_fee = get_fee_rate(self.og1.market_type, self.og1.exchange)
    assert abs(self._pass_maker_fee) <= 1e-2, self._pass_maker_fee
    assert abs(self._agg_taker_fee) <= 1e-2, self._agg_taker_fee
    self.logger.info('%s init pos0 %f, pos1 %f, expected agg pos %f' %
                     (self._ticker0, self._pos0, self._pos1, self._expected_agg_pos))

  def _get_sell_price_pull(self, price):
    if price is not None:
      return price * (1 - self._pull_edge)
    return None

  def _get_buy_price_pull(self, price):
    if price is not None:
      return price * (1 + self._pull_edge)
    return None

  def _get_sell_price_push(self, price):
    if price is not None:
      return price * (1 + self._push_edge)
    return None

  def _get_buy_price_push(self, price):
    if price is not None:
      return price * (1 - self._push_edge)
    return None

  def _check_buy_price_within(self, buy_price, bids):
    if self._post_pass_within_levels > 0 and \
            bids[:self._post_pass_within_levels][-1][0] > buy_price:
      return False
    if self._post_pass_within_lotsize > 0.01:
      cum_amt = 0.0
      for bid in bids:
        if bid[0] < buy_price:
          break
        cum_amt += bid[1]
      if cum_amt > self._post_pass_within_lotsize * self._lot_size:
        return False
    return True

  def _check_sell_price_within(self, sell_price, asks):
    if self._post_pass_within_levels > 0 and \
            asks[:self._post_pass_within_levels][-1][0] < sell_price:
      return False
    if self._post_pass_within_lotsize > 0.01:
      cum_amt = 0.0
      for ask in asks:
        if ask[0] > sell_price:
          break
        cum_amt += ask[1]
      if cum_amt > self._post_pass_within_lotsize * self._lot_size:
        return False
    return True

  def _get_pass_buy_price(self, price, bids, asks, tick):
    price = tick.floor(price)
    if price >= asks[0][0]:
      price = tick.get_prev_price(asks[0][0])
    if self._check_post_within:
      if not self._check_buy_price_within(price, bids) \
              and self._last_pass_buy_price < price:
        price = None
      if price is not None:
        self._last_pass_buy_price = price
    return price

  def _get_pass_sell_price(self, price, bids, asks, tick):
    price = tick.ceil(price)
    if price <= bids[0][0]:
      price = tick.get_next_price(bids[0][0])
    if self._check_post_within:
      if not self._check_sell_price_within(price, asks) \
              and self._last_pass_sell_price > price:
        price = None
      if price is not None:
        self._last_pass_sell_price = price
    return price

  def _get_agg_working_qty(self):
    # Use working orders to update |self._expected_working_qty|.
    if self._get_ts() - self._last_agg_submit_fill_ts > 1 * 10**9:
      order_id_set = set()
      working_qty = 0.0
      for order in self._agg_executor1._wo.buy_working_orders:
        working_qty += (order.qty - self._agg_working_order_fills.get(order.order_id, 0.0))
        order_id_set.add(order.order_id)
      for order in self._agg_executor1._wo.sell_working_orders:
        working_qty -= (order.qty - self._agg_working_order_fills.get(order.order_id, 0.0))
        order_id_set.add(order.order_id)
      # Remove order ids that are not in working orders.
      obsolete_ids = self._agg_working_order_fills.keys() - order_id_set
      for order_id in obsolete_ids:
        del self._agg_working_order_fills[order_id]
      self.logger.info('working qty expected %f, updated %f' %
                       (self._expected_working_qty, working_qty))
      self._expected_working_qty = working_qty
    return self._expected_working_qty

  def _get_agg_qty(self):
    lot_size = self._config1['lot_size']
    working_qty = self._get_agg_working_qty()
    agg_balance = self._get_agg_balance()
    imbalance = self._expected_agg_pos - agg_balance - working_qty
    self.logger.info('lot_size: %f, expected agg position: %f, '
                     'agg position: %f, working qty: %f, imbalance: %f' %
                     (lot_size, self._expected_agg_pos, agg_balance, working_qty, imbalance))
    imbalance = min(imbalance, lot_size)
    imbalance = max(imbalance, -lot_size)
    return imbalance

  def _position_equal(self, p1, p2):
    return abs(p1 - p2) < self._equal_theshold

  def _guess_filled(self, bids, asks):
    new_pass_pos = self.og0.get_balance(self._ticker0)
    buy_orders = self._pass_executor0._wo.buy_working_orders
    sell_orders = self._pass_executor0._wo.sell_working_orders
    if not self._position_equal(new_pass_pos, self._latest_pass_pos):
      fill_qty = 0.0
      fill_qty = new_pass_pos - self._latest_pass_pos
      prev_pos = self._latest_pass_pos
      if fill_qty > 0:
        if new_pass_pos > self._pos0 > self._latest_pass_pos:
          fill_qty = new_pass_pos - self._pos0
          self.logger.info('Concurrent fill msg and balance update, '
                           'new pos %f, last pos %f, pos0 %f' %
                           (new_pass_pos, self._latest_pass_pos, self._pos0))
        self._latest_pass_pos = new_pass_pos
        self.on_passive_fill(None, fill_qty, None, True)
      elif fill_qty < 0:
        if new_pass_pos < self._pos0 < self._latest_pass_pos:
          fill_qty = new_pass_pos - self._pos0
          self.logger.info('Concurrent fill msg and balance update, '
                           'new pos %f, last pos %f, pos0 %f' %
                           (new_pass_pos, self._latest_pass_pos, self._pos0))
        self._latest_pass_pos = new_pass_pos
        self.on_passive_fill(None, fill_qty, None, True)
      self.logger.info(
          'Balance guess %s pass filled, prev pos: %f, new pos: %f, '
          'fill_qty: %f, pos from fill %f' %
          (self._ticker0, prev_pos, new_pass_pos, fill_qty, self._latest_pass_pos_from_fill))

    # Fast but not reliable.
    if self._guess_fill_from_book:
      if len(buy_orders) > 0 and bids[0][0] < buy_orders[0].price:
        self.logger.info('Order %s %s guess buy filled, price %f, qty %f, bid0 %f' %
                         (self._ticker0,
                          buy_orders[0].order_id,
                          buy_orders[0].price,
                          buy_orders[0].qty,
                          bids[0][0]))
      if len(sell_orders) > 0 and asks[0][0] > sell_orders[0].price:
        self.logger.info('Order %s %s guess sell filled, price %f, qty %f, ask0 %f' %
                         (self._ticker0,
                          sell_orders[0].order_id,
                          sell_orders[0].price,
                          sell_orders[0].qty,
                          asks[0][0]))

  def _get_book_invalid_prices(self, bids, asks):
    return bids[0][0] * 0.95, asks[0][0] * 1.05

  def on_book(self, exchange_id, book, mean_scale, spread_std):
    if self._pass_executor0 and self._agg_executor1:
      self._book_ts[exchange_id] = book.timestamp
      book_key0 = (self._product0.exchange, self._ticker0)
      book_key1 = (self._product1.exchange, self._ticker1)
      bids0, asks0 = self.feed[book_key0]
      bids1, asks1 = self.feed[book_key1]

      if not self._is_book_valid():
        if self._get_ts() - self._last_feed_invalid_log_ts > 2e9:
          self.logger.warning('%s book is not updated %s %s.' %
                              (self._ticker0, self._book_ts[0], self._book_ts[1]))
          self._last_feed_invalid_log_ts = self._get_ts()
        invalid_buy_price, invalid_sell_price = \
            self._get_book_invalid_prices(bids0, asks0)
        # Prevent fills by adjusting prices.
        self._pass_executor0.manage_mm_orders(invalid_sell_price, invalid_buy_price)
        return

      if not FLAGS.dry_run:
        if self._guess_fill_from_balance:
          self._guess_filled(bids0, asks0)
        pass_buy_price0, pass_sell_price0 = self._pricer.calculate_price(
            bids0, asks0, bids1, asks1, mean_scale, spread_std, self._threshold, self._lot_size)
        pass_buy_price0 *= (1 + self._parent.pos_lean)
        pass_sell_price0 *= (1 + self._parent.pos_lean)
        pass_buy_price0 = self._get_pass_buy_price(pass_buy_price0,
                                                   bids0,
                                                   asks0,
                                                   self._tick_info0[0])
        pass_sell_price0 = self._get_pass_sell_price(pass_sell_price0,
                                                     bids0,
                                                     asks0,
                                                     self._tick_info0[0])
        pass_buy_qty, pass_sell_qty = self._pricer.calculate_qty(bids1, asks1,
                                                                 self._get_pass_balance(),
                                                                 self._max_pos, self._min_pos,
                                                                 self._min_pass_order_qty,
                                                                 self._lot_size)
        self._pass_executor0.manage_mm_orders(
            pass_sell_price0,
            pass_buy_price0,
            sell_price_pull=self._get_sell_price_pull(pass_sell_price0),
            buy_price_pull=self._get_buy_price_pull(pass_buy_price0),
            sell_price_push=self._get_sell_price_push(pass_sell_price0),
            buy_price_push=self._get_buy_price_push(pass_buy_price0),
            buy_qty=pass_buy_qty,
            sell_qty=pass_sell_qty)
        self._agg_executor1.manage_agg_orders()
      self._spread_std = spread_std

  def _get_pass_fillinfo(self, fill_price, fill_qty, order):
    return FillInfo(action='pass filled',
                    symbol=self._product0.symbol,
                    exchange=self._product0.exchange,
                    fill_price=fill_price,
                    fill_qty=fill_qty,
                    order_side=order.order_side,
                    timestamp=self._get_ts())

  def _is_fill_msg_delayed(self, order_side, signed_qty):
    if order_side == OrderSide.BUY:
      return signed_qty <= 0
    else:
      return signed_qty >= 0

  def _is_balance_update_delayed(self, fill_qty):
    return abs(self._latest_pass_pos_from_fill + fill_qty - self._latest_pass_pos) > \
        abs(self._latest_pass_pos_from_fill - self._latest_pass_pos)

  def on_passive_fill(self, fill_price, fill_qty, order, is_guess=False):
    self._fill0_ts = self._get_ts()
    if not is_guess:
      self._fillinfo_writer.write(self._get_pass_fillinfo(fill_price, fill_qty, order))
      signed_qty = fill_qty
      if order.order_side == OrderSide.SELL:
        signed_qty = -signed_qty
      else:
        signed_qty = signed_qty * (1 - self._pass_maker_fee)
      latest_pass_pos_from_fill = self._latest_pass_pos_from_fill + signed_qty
      signed_qty = latest_pass_pos_from_fill - self._pos0
      self.logger.info('%s fill, balance pos %f, fill pos %f, pos0 %f, fill qty %f, '
                       'signed qty %f' % (self._product0.exchange,
                                          self._latest_pass_pos,
                                          self._latest_pass_pos_from_fill,
                                          self._pos0,
                                          fill_qty,
                                          signed_qty))
      if (self._guess_fill_from_balance
          and self._is_fill_msg_delayed(order.order_side, signed_qty)) \
              or self._guess_fill_from_balance_only:
        # Balance fill already came.
        self._latest_pass_pos_from_fill = latest_pass_pos_from_fill
        self.logger.info('%s fill msg is delayed.' % self._product0.exchange)
        if self._position_equal(self._latest_pass_pos_from_fill, self._latest_pass_pos):
          self._latest_pass_pos_from_fill = self._latest_pass_pos
        return
      self._latest_pass_pos_from_fill = latest_pass_pos_from_fill
      fill_qty = signed_qty
    else:
      if (not self._guess_fill_from_balance_only) and \
              self._is_balance_update_delayed(fill_qty):
        self.logger.info('%s balance update delayed.' % self._product0.exchange)
        return

    # Aggressive on product1
    product = self._product1
    prices = self.feed[(product.exchange, self._ticker1)]
    if fill_qty < 0:
      # Pass sell
      fill_qty = -fill_qty
      agg_order_side = OrderSide.BUY
      price = self._tick_info1[0].get_next_price(prices[1][0][0] * (1 + self._aggressive_amt))
      self._expected_agg_pos += fill_qty
      self._update_pos0(self._pos0 - fill_qty)
    else:
      # Pass buy
      agg_order_side = OrderSide.SELL
      price = self._tick_info1[0].get_prev_price(prices[0][0][0] * (1 - self._aggressive_amt))
      self._expected_agg_pos -= fill_qty
      self._update_pos0(self._pos0 + fill_qty)
    # Cannot trust pass balance.
    qty = self._get_agg_qty()

    submitted = False
    if (agg_order_side == OrderSide.BUY and qty > 0) or \
            (agg_order_side == OrderSide.SELL and qty < 0):
      submitted = True
      agg_order_qty = qty
      if agg_order_qty < 0:
        agg_order_qty = self._tick_info1[1].floor(-agg_order_qty)
      ret = self._agg_executor1.submit_agg_order(order_side=agg_order_side,
                                                 order_price=price,
                                                 order_qty=agg_order_qty)
      if ret:
        self._expected_working_qty += qty
        self._last_agg_submit_fill_ts = self._get_ts()
      self.logger.info('agg submitted return code %s' % ret)
    if submitted:
      self._fillinfo_writer.write(
          FillInfo(action='agg submitted',
                   symbol=self._product1.symbol,
                   exchange=self._product1.exchange,
                   fill_price=price,
                   fill_qty=agg_order_qty,
                   order_side=agg_order_side,
                   timestamp=self._get_ts()))
      self.logger.info('%s %s %s %s' %
                       (self._ticker1, prices[0][0], prices[1][0], self._book_ts[1]))
    self.logger.info('expected %s positions, %f %f, fill qty: %f, qty: %f, std: %f' %
                     (self._ticker0, self._pos0, self._pos1, fill_qty, qty, self._spread_std))

  def _update_expected_agg_pos(self):
    self._expected_agg_pos = self._config0['max_pos'] - self._pos0

  def on_aggressive_fill(self, fill_price, fill_qty, order):
    # Cannot trust agg balance.
    self._update_expected_agg_pos()
    self._agg_working_order_fills[order.order_id] = \
        self._agg_working_order_fills.get(order.order_id, 0.0) + fill_qty
    self.logger.info('expected agg pos: %f' % (self._expected_agg_pos))
    self._fillinfo_writer.write(
        FillInfo(action='agg filled',
                 symbol=order.product.symbol,
                 exchange=order.product.exchange,
                 fill_price=fill_price,
                 fill_qty=fill_qty,
                 order_side=order.order_side,
                 timestamp=self._get_ts()))
    if order.order_side == OrderSide.BUY:
      self._update_pos1(self._pos1 + fill_qty * (1 - self._agg_taker_fee))
      self._expected_working_qty -= fill_qty
    else:
      self._update_pos1(self._pos1 - fill_qty)
      self._expected_working_qty += fill_qty
    self._last_agg_submit_fill_ts = self._get_ts()
    self.logger.info('expected %s pos, %f %f' % (self._ticker0, self._pos0, self._pos1))

  def on_trade(self, exchange_id, trade):
    pass


class SpreadBasisStrategy(object):
  def __init__(self, strategy_params, trading_tickers):
    self.logger = logging.getLogger(__name__)

    self._trading_tickers = set(trading_tickers)
    self.logger.info('Traded tickers %s' % self._trading_tickers)
    self._config = strategy_params['config']
    self._threshold = self._config['common']['threshold']
    self._pull_edge = self._config['common']['pull_edge_bp']
    self._push_edge = self._config['common']['push_edge_bp']
    self._ema_weight = self._config['common']['ema_weight']
    self._min_basis_std = self._config['common']['minimum_basis_std']
    self._max_basis_std = self._config['common'].get('maximum_basis_std', 1e3)
    self._volatile_bases = self._config['volatile_configs']['bases']
    self._volatile_threshold = self._config['volatile_configs']['theshold']
    self._max_pos_lean = FLAGS.max_pos_lean_bp * 1e-4
    self.pos_lean = 0.0

    self._exch0_subreq = SubscriptionRequest.from_dict(self._config['common']['exch0_subreq'])
    self._exch1_subreq = SubscriptionRequest.from_dict(self._config['common']['exch1_subreq'])

    self._exchange0 = self._exch0_subreq.exchange
    self._exchange1 = self._exch1_subreq.exchange
    self._og_config = strategy_params['og_config']

    self._balance_checker = None
    account_checker_config = self._get_account_checker_config()
    if account_checker_config is not None:
      self._balance_checker = AccountChecker(account_checker_config)

    self._book_builder_clses = (get_flow_book_builder_cls(self._exch0_subreq),
                                get_flow_book_builder_cls(self._exch1_subreq))
    self._live_env_clses = (get_live_env_cls(self._exch0_subreq),
                            get_live_env_cls(self._exch1_subreq))
    self._live_env0 = None
    self._live_env1 = None

    self.og0 = None
    self.og1 = None

    self._products0, self._products1 = _gen_arb_products(self._config['common'],
                                                         self._config['products0'],
                                                         self._config['products1'],
                                                         self._book_builder_clses[0].ProductType,
                                                         self._book_builder_clses[1].ProductType)

    exch0_exchange_feed_checker = self._config['common'].get('exch0_exchange_feed_checker', False)
    exch1_exchange_feed_checker = self._config['common'].get('exch1_exchange_feed_checker', False)
    feed_checker_products = []
    if exch0_exchange_feed_checker:
      feed_checker_products += self._products0.products
    if exch1_exchange_feed_checker:
      feed_checker_products += self._products1.products
    feed_checker_config = get_default_feed_checker_config(feed_checker_products)

    self._feed = MultiFeedSubsystem(feed_checker_config=feed_checker_config)
    self._ts = None

    self._order_canceled_after_og_ready = False
    self._exit_attempt = 0
    self._exit_executor = None
    self._balance_printer = None
    self._fillinfo_writer = FillInfoWriter('fills.csv', FLAGS.strategy_name)

    self._book_levels = FLAGS.book_levels
    self._last_feed = {}
    self._std_diff = np.nan
    self._mean_scale = np.nan
    self._price_diffs = None
    tickers = set()
    for product in self._products0.products:
      tickers.add(_get_ticker(product))

    self._tickers = np.array(sorted(list(tickers)))
    self._tickers_map = {ticker: i for i, ticker in enumerate(self._tickers)}
    self._volatile_mask = np.isin(self._tickers, self._volatile_bases)
    self._exchanges = [self._exchange0, self._exchange1]
    self._last_price = np.empty((2, len(self._tickers)))
    self._last_pos = np.empty((2, len(self._tickers)))
    self._last_price[:] = np.nan
    self._last_pos[:] = np.nan
    self._max_pos = self._get_max_pos()

    self._reserve_map = self._get_reserve_map()
    self._strategy_request = convert_strategy_request_from_dict(
        self._config['common']['strategy_request'])
    self._reserve_proto0 = reserve_as_proto(self._exch0_subreq, self._reserve_map)
    self._reserve_proto1 = reserve_as_proto(self._exch1_subreq, self._reserve_map)
    self._strategy_logger = StrategyLogProvider(self._strategy_request)

    self._substrategies = {}
    for product0, product1 in zip(self._products0.products, self._products1.products):
      ticker = _get_ticker(product0)
      threshold = self._volatile_threshold if ticker in self._volatile_bases \
          else self._threshold
      self._substrategies[ticker] = SingleSpreadBasisSubStrategy(
          self,
          product0,
          product1,
          self._pull_edge,
          self._push_edge,
          trading_threshold=threshold,
          fillinfo_writer=self._fillinfo_writer)
    self._dump_pnl_balance = FLAGS.dump_pnl_balance

  @staticmethod
  def _get_config_key(product):
    key = product.subscription_symbol
    if product.exchange == 'Huobi' and product.product_type == 'Futures':
      key = key.replace('CURRENT', 'THIS')
      key = key.replace('THIS_QUARTER', 'QUARTER')
    return key

  def _get_max_pos(self):
    max_pos = np.empty((len(self._tickers),))
    for i, product in enumerate(self._products0.products):
      key = SpreadBasisStrategy._get_config_key(product)
      max_pos[i] = self._config['products0'][key]['max_pos']
    return max_pos

  def _get_reserve_map(self):
    reserves = dict()
    for product in self._products0.products:
      key = SpreadBasisStrategy._get_config_key(product)
      reserves[product.base.currency] = (self._config['products0'][key]['max_pos']
                                         + self._config['products0'][key]['min_pos']) / 2.0
    return reserves

  def periodic_write_logs(self):
    self.og0.og_logger.write_account_request()
    self.og1.og_logger.write_account_request()
    self._strategy_logger.write_request()
    self._strategy_logger.write_reserve(self._reserve_proto0,
                                        account_request=self.og0.og_logger.request)
    self._strategy_logger.write_reserve(self._reserve_proto1,
                                        account_request=self.og1.og_logger.request)
    self._ioloop.add_timeout(datetime.timedelta(hours=1), self.periodic_write_logs)

  def _get_account_checker_config(self):
    if self._config.get('account_checker') is None:
      return None

    config_keys = {
        'reserve',
        'lot_size',
        'stack_size',
        'ratio_min_balance',
        'ratio_max_balance',
        'ratio_min_hold_balance',
    }

    balance_check_config_list = []
    for currency, config in self._config['account_checker'].items():
      config = {**config}
      if 'lot_size' not in config:
        for product in self._config['products']:
          base = SharedProduct.FromStrNativeProduct(product).base.currency
          if base == currency:
            config['lot_size'] = self._config['products'][product]['lot_size']
            break

      assert config_keys == set(config.keys()), config
      reserve = config['reserve']
      config['currency'] = currency
      config['min_balance'] = reserve * config['ratio_min_balance']
      config['max_balance'] = reserve * config['ratio_max_balance']
      config['min_available_balance'] = 0.0
      config['max_available_balance'] = config['lot_size'] * config['stack_size']
      config['min_hold_balance'] = reserve * config['ratio_min_hold_balance']
      config['max_hold_balance'] = reserve * config['ratio_max_balance']
      currency_checker_config = convert_checker_config_from_dict(config)
      balance_check_config_list.append(currency_checker_config)

    return AccountCheckerConfig(each_balance_check_config=balance_check_config_list)

  def print_exchange_balance(self):
    if self.og0 and self.og1:
      quote0 = self.og0.get_balance(self._products0.quote_currency)
      quote1 = self.og1.get_balance(self._products1.quote_currency)
      pos0 = []
      pos1 = []
      for prod0, prod1 in zip(self._products0.products, self._products1.products):
        pos0.append(self.og0.get_balance(prod0.base))
        pos1.append(self.og1.get_balance(prod1.base))
      requires0 = np.sum(self._last_price[0] * self._max_pos)
      requires1 = np.sum(self._last_price[1] * self._max_pos)
      quote0 += np.sum(self._last_price[0] * np.array(pos0))
      quote1 += np.sum(self._last_price[1] * np.array(pos1))
      self.logger.info('Extra balance0: %s, balance1: %s' %
                       (quote0 - requires0, quote1 - requires1))

  def _update_prices(self, exchange_id, product, book):
    ticker = _get_ticker(product)
    book_key = (product.exchange, ticker)
    bids = book.get_bid_array(self._book_levels)
    asks = book.get_ask_array(self._book_levels)
    if len(bids) == 0 or len(asks) == 0:
      return False
    self._last_feed[book_key] = (bids, asks)
    tid = self._tickers_map[ticker]
    self._last_price[exchange_id][tid] = SpreadBasisStrategy.get_mid(bids, asks)
    self._last_pos[0][tid] = self._substrategies[ticker]._pos0
    self._last_pos[1][tid] = self._substrategies[ticker]._pos1
    cashpos = np.sum(self._last_price * self._last_pos, 1)
    # Postive means exchange0 has less coin balance, shift passive price higher.
    cashpos[1] *= self._mean_scale
    self.pos_lean = np.nan_to_num((cashpos[1] - cashpos[0]) /
                                  (cashpos[0] + cashpos[1]) * self._max_pos_lean)
    return True

  @staticmethod
  def get_mid(bids, asks):
    return (bids[0][0] + asks[0][0]) / 2.0

  @staticmethod
  def _get_ema(old, new, weight):
    if not np.isnan(old):
      return old * (1 - weight) + new * weight
    return new

  def _update_stats(self, p0, p1):
    implied_scale = 1.0 * p0 / p1
    mean_scale = np.mean(implied_scale[~self._volatile_mask])
    self._mean_scale = self._get_ema(self._mean_scale, mean_scale, self._ema_weight)
    mid = (p0 + p1 * self._mean_scale) / 2.0
    self._price_diffs = (p0 - p1 * self._mean_scale) / mid
    std = np.std(self._price_diffs[~self._volatile_mask])
    self._std_diff = self._get_ema(self._std_diff, std, self._ema_weight)

  def on_book(self, product, exchange_id, book):
    self._ts = book.timestamp
    if self._update_prices(exchange_id, product, book):
      p0 = self._last_price[0]
      p1 = self._last_price[1]
      self._update_stats(p0, p1)
      for ticker, strat in self._substrategies.items():
        if ticker in self._trading_tickers:
          strat.on_book(exchange_id,
                        book,
                        self._mean_scale,
                        min(max(self._std_diff, self._min_basis_std), self._max_basis_std))
      if self._balance_printer:
        self._balance_printer.on_book(product, exchange_id, book)

  def on_trade(self, exchange_id, trade):
    for _, strat in self._substrategies.items():
      strat.on_trade(exchange_id, trade)

  def on_passive_fill(self, fill_price, fill_qty, order):
    ticker = _get_ticker(order)
    if ticker in self._trading_tickers:
      self._substrategies[ticker].on_passive_fill(fill_price, fill_qty, order)

  def on_aggressive_fill(self, fill_price, fill_qty, order):
    ticker = _get_ticker(order)
    if ticker in self._trading_tickers:
      self._substrategies[ticker].on_aggressive_fill(fill_price, fill_qty, order)

  def signal_handler(self, sig, frame):
    self.logger.warning('Caught signal: %s', sig)
    self._ioloop.add_callback_from_signal(self.on_exit)

  def on_exit(self):
    self._exit_attempt += 1
    if self._exit_attempt > 5:
      self._ioloop.add_callback(self._ioloop.stop)

    finished = True
    if self._exit_executor is not None and not FLAGS.dry_run:
      ogs = to_list([self.og0, self.og1])
      self.logger.info('Exit callback executed.')
      finished = self._exit_executor.cancel_all_on_exit(ogs)

    if finished:
      self._ioloop.add_callback(self._ioloop.stop)
    else:
      timeout = datetime.timedelta(seconds=2)
      self._ioloop.add_timeout(timeout, self.on_exit)

  def start(self, ioloop, queue, flow_sub):
    self._ioloop = ioloop
    self._feed.start(ioloop, queue, flow_sub)
    self._live_env0.start()
    self._live_env1.start()
    self.dump_pnl_balance(True)

  def subscribe_feeds(self):
    for product in self._products0.products:
      self._feed.subscribe_book([product], functools.partial(self.on_book, product, 0))
    for product in self._products1.products:
      self._feed.subscribe_book([product], functools.partial(self.on_book, product, 1))
    if FLAGS.subscribe_trade:
      self._feed.subscribe_trade(self._products0.products, functools.partial(self.on_trade, 0))
      self._feed.subscribe_trade(self._products1.products, functools.partial(self.on_trade, 1))

  def subscribe_live_envs(self, ioloop, queue, flow_sub):
    og_params = self._config.get('og_params', {})
    og_params0 = convert_og_config_from_dict(self._exch0_subreq, og_params)
    og_params1 = convert_og_config_from_dict(self._exch1_subreq, og_params)
    self._live_env0 = create_live_env_instance(self._live_env_clses[0],
                                               self._products0.currencies,
                                               self._products0.products,
                                               self._og_config[0],
                                               og_params=og_params0,
                                               ioloop=ioloop,
                                               queue=queue,
                                               flow_sub=flow_sub,
                                               og_log_callback=self._strategy_logger.write_og_log)
    self._live_env0.on_order_gateway_reset_callback.append(self.register_order_gateway)

    self._live_env1 = create_live_env_instance(self._live_env_clses[1],
                                               self._products1.currencies,
                                               self._products1.products,
                                               self._og_config[1],
                                               og_params=og_params1,
                                               ioloop=ioloop,
                                               queue=queue,
                                               flow_sub=flow_sub,
                                               og_log_callback=self._strategy_logger.write_og_log)
    self._live_env1.on_order_gateway_reset_callback.append(self.register_order_gateway)

  def register_balance_printer(self):
    ogs = [self.og0, self.og1]
    base_currency_reserves = {
        currency: self._config['products0'][self._products0.products[i].symbol]['max_pos'] for i,
        currency in enumerate(self._products0.base_currencies)
    }
    self._balance_printer = KimchiEstimatedBalancePrinter(
        quote_currencies=(self._products0.quote_currencies + self._products1.quote_currencies),
        base_currencies=self._products0.base_currencies,
        base_currency_reserves=base_currency_reserves,
        ogs=ogs)
    self.dump_pnl_balance()

  def dump_pnl_balance(self, periodic=False):
    self.print_balance()
    self.print_exchange_balance()
    if self._balance_printer is not None and self._dump_pnl_balance:
      pnl_balances = self._balance_printer.calculate_pnl_balance()
      if pnl_balances is not None:
        for pnl_balance in pnl_balances:
          self._strategy_logger.write_pnl_balance(pnl_balance=pnl_balance,
                                                  quote=None,
                                                  subreq_id=None)
    if periodic:
      self._ioloop.add_timeout(datetime.timedelta(seconds=20), self.dump_pnl_balance, True)

  def print_balance(self):
    if self._balance_printer is not None:
      imbalance, quote0, quote1 = \
          self._balance_printer.print_balance(self._mean_scale)
      self.logger.info('std_diff %f, price_diffs %s, pos_lean %s' %
                       (self._std_diff, self._price_diffs, self.pos_lean))

  def register_order_gateway(self, name, order_gateway):
    self.logger.info('OG initialized: %s', name)
    if order_gateway.exchange == self._exchange0:
      self.og0 = order_gateway
    elif order_gateway.exchange == self._exchange1:
      self.og1 = order_gateway
    # cancel working orders first after OGs are ready.
    if not self._order_canceled_after_og_ready:
      order_gateway.cancel_all()
      self.logger.info('Orders canceled after OGs are ready.')
      self._order_canceled_after_og_ready = True
    self._exit_executor = SimpleCancelAllExecutor()

    if self.og0 and self.og1:
      self.register_balance_printer()
      self.og0.register_fill_callback(self.on_passive_fill)
      self.og1.register_fill_callback(self.on_aggressive_fill)
      for prod0, prod1 in zip(self._products0.products, self._products1.products):
        ticker = _get_ticker(prod0)
        config0 = self._config['products0'][prod0.symbol].copy()
        config1 = self._config['products1'][prod1.symbol].copy()
        config0.update(self._config.get('pass_executor_configs', {}))
        config1.update(self._config.get('agg_executor_configs', {}))
        self._substrategies[ticker].init_executors(self._exch0_subreq,
                                                   self._exch1_subreq,
                                                   config0,
                                                   config1,
                                                   self.og0,
                                                   self.og1)
      if self._dump_pnl_balance:
        self.periodic_write_logs()


def get_param_set(og_configs, config_filename):
  og_config = og_configs.split(',')
  config = load_clean_json(config_filename)
  return {
      'og_config': og_config,
      'config': config,
  }


def main(argv):
  assert FLAGS.strategy_name

  strategy_params = get_param_set(og_configs=FLAGS.og_configs,
                                  config_filename=FLAGS.config_filename)

  queue = SimpleQueue('spreadbasis')
  flow_sub = flow_subscriber.from_queue(queue)
  ioloop = IOLoop.current()

  trading_tickers = FLAGS.trading_tickers.split(',')
  strategy = SpreadBasisStrategy(strategy_params, trading_tickers)
  strategy.subscribe_feeds()
  strategy.subscribe_live_envs(ioloop, queue, flow_sub)
  strategy.start(ioloop, queue, flow_sub)

  exit_after = datetime.timedelta(hours=6)
  ioloop.add_timeout(exit_after, strategy.on_exit)

  signal.signal(signal.SIGTERM, strategy.signal_handler)
  signal.signal(signal.SIGINT, strategy.signal_handler)

  ioloop.start()


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
